package com.weijian.service.impl;

import cn.hutool.core.lang.UUID;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.weijian.advisor.MyLoggerAdvisor;
import com.weijian.advisor.RagLoggerAdvisor;
import com.weijian.common.MessageWrapper;
import com.weijian.common.pojo.CommonResult;
import com.weijian.dal.dataobject.BaseConversation;
import com.weijian.dal.vo.ConversationVO;
import com.weijian.service.ChatService;
import com.weijian.service.MpConversationService;
import com.weijian.utils.ParseWordFileUtil;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.util.List;

import static com.weijian.constants.SystemPrompt.DEEPPROMPT;
import static com.weijian.constants.SystemPrompt.USUALPROMPT;

@Service
@RequiredArgsConstructor
public class ChatServiceImpl implements ChatService {

    final ChatMemory chatMemory;

    final ChatClient chatClient;

    final VectorStore vectorStore;

    final MpConversationService conversationService;


    @Override
    public ConversationVO create() {
        String conversationId = UUID.randomUUID().toString();
        String title = "conversation-" + conversationId;
        conversationService.save(BaseConversation.builder()
                .title(title)
                .conversationId(conversationId)
                .build());
        return ConversationVO.builder().conversationId(conversationId)
                .build();
    }

    @Override
    public void edit(String conversationId, String title) {
        if (!StringUtils.hasText(title)) {
            return;
        }
        conversationService.update(BaseConversation.builder()
                        .title(title)
                        .build(),
                new LambdaQueryWrapper<BaseConversation>().eq(BaseConversation::getConversationId, conversationId));
    }

    @Override
    public List<ConversationVO> list() {
        return conversationService.list(new LambdaQueryWrapper<BaseConversation>().orderByAsc(BaseConversation::getUpdateTime))
                .stream().map(v -> ConversationVO.builder()
                        .conversationId(v.getConversationId())
                        .build()).toList();
    }

    @Override
    public void delete(String conversationId) {
        conversationService.remove(
                new LambdaQueryWrapper<BaseConversation>().eq(BaseConversation::getConversationId, conversationId));
    }

    @Override
    public ConversationVO get(String conversationId) {
        return conversationService.getOneOpt(
                        new LambdaQueryWrapper<BaseConversation>().eq(BaseConversation::getConversationId, conversationId))
                .map(v -> ConversationVO.builder().conversationId(v.getConversationId())
                        .title(v.getTitle())
                        .content(MessageWrapper.fromConversationStr(v.getContent()))
                        .build()
                ).orElseThrow(() -> new RuntimeException("can not find conversation with id: %s ".formatted(conversationId)));

    }

    @Override
    public Flux<String> chat(String conversationId, String message, List<MultipartFile> files, String tagNames) {
        if (files == null || files.isEmpty()) {
            return textChat(USUALPROMPT, message, conversationId, tagNames);
        } else {
            return multiModalChatAndStore(USUALPROMPT, message, conversationId, files, tagNames);
        }
    }

    @Override
    public Flux<String> deepChat(String conversationId, String message, List<MultipartFile> files, String tagNames) {
        if (files == null || files.isEmpty()) {
            return textChat(DEEPPROMPT, message, conversationId, tagNames);
        } else {
            return multiModalChatAndStore(DEEPPROMPT, message, conversationId, files, tagNames);
        }
    }

    @Override
    public CommonResult<String> search(String message, String title) {
        DocumentRetriever retriever = VectorStoreDocumentRetriever.builder()
                .vectorStore(vectorStore)
                .similarityThreshold(0.73)
                .topK(5)
                .filterExpression(title != null && !title.isEmpty()
                        ? () -> new FilterExpressionBuilder().eq("title", title).build()
                        : () -> null)
                .build();
        List<Document> documents = retriever.retrieve(new Query(message));
        assert documents != null;
        List<Object> idList = documents.stream().map(document -> document.getMetadata().get("title")).toList();
        return CommonResult.success(idList.toString());
    }


    @SneakyThrows
    private Flux<String> multiModalChatAndStore(String prompt, String message, String conversationId,
                                                List<MultipartFile> files, String title) {
        MessageChatMemoryAdvisor chatMemoryAdvisor = MessageChatMemoryAdvisor.builder(chatMemory).conversationId(conversationId).build();
        Advisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .similarityThreshold(0.50)
                        .topK(3)
                        .vectorStore(vectorStore)
                        .filterExpression(title != null && !title.isEmpty()
                                ? () -> new FilterExpressionBuilder().eq("title", title).build()
                                : () -> null)
                        .build())
                .build();
        return Flux.fromIterable(files)
                .flatMap(file -> {
                    try {
                        String fileName = file.getOriginalFilename();
                        String content = "";
                        if (fileName != null && fileName.endsWith(".docx")) {
                            content = ParseWordFileUtil.parseDocxFile(file);
                        } else if (fileName != null && fileName.endsWith(".doc")) {
                            content = ParseWordFileUtil.parseDocFile(file);
                        } else {
                            return Flux.error(new IllegalArgumentException("文件类型不支持，请上传word文件: " + fileName));
                        }
                        return Flux.just(content);
                    } catch (Exception e) {
                        return Flux.error(e);
                    }
                })
                .flatMap(fileContent -> chatClient.prompt(prompt)
                        .user(p -> p.text(message).text(fileContent))
                        .advisors(retrievalAugmentationAdvisor)
                        .advisors(chatMemoryAdvisor)
                        .stream()
                        .content());
    }


    private Flux<String> textChat(String prompt, String message, String conversationId, String title) {
        Advisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .similarityThreshold(0.50)
                        .topK(3)
                        .vectorStore(vectorStore)
                        .filterExpression(title != null && !title.isEmpty()
                                ? () -> new FilterExpressionBuilder().eq("title", title).build()
                                : () -> null)
                        .build())
                .build();
        MessageChatMemoryAdvisor chatMemoryAdvisor = MessageChatMemoryAdvisor.builder(chatMemory).conversationId(conversationId).build();
        return chatClient.prompt(prompt)
                .user(message)
                .advisors(retrievalAugmentationAdvisor)
                .advisors(chatMemoryAdvisor)
                .advisors(new MyLoggerAdvisor())
                .advisors(new RagLoggerAdvisor(vectorStore,chatClient))
                .stream()
                .content();
    }

}
